# -*- coding: utf-8 -*-

'''
 :author: Rayco Abad-Martín <rayco.abad@gmail.com>
 :author: Elena Marrero-Méndez <elemarez@gmail.com>
'''

from yacf.Backends.C.Writers.CWriter import CWriter
from yacf.Backends.Common.Mutators.AbstractMutator import IgnoreMutationException, AbstractMutator, AbortMutationException
from yacf.Backends.Common.Mutators.AstSupport import IDNameMutator
from yacf.Backends.Common.Visitors.GenericVisitors import *
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.MiddleEnd.Functions.Mutators import Outliner
from yacf.MiddleEnd.Functions.Mutators.Outliner import block_is_SESE, OutlineMutator, ParameterStyle
from yacf.MiddleEnd.Loop.Analysis import OmpParametrize
from yacf.Tools import SourceStorage
from yacf.Tools.Debug import DotDebugTool
from yacf.Tools.Declarations import type_of_id, decl_of_id
from yacf.Tools.Tree import InsertTool, NodeNotFound, ReplaceTool, RemoveTool
import cStringIO
import config

c_ast = getCurrentLanguageAst()

class MPIMutatorError(Exception):
    def __init__(self, description):
        self.description = description

    def __str__(self):
        return "MPIMutatorError :: " + self.description

class AbstractMPIMutator(Outliner.OutlineMutator):
    """ Common methods to work with MPI
    """
    def __init__(self, clauses = None, source_storage = None):
        if clauses == None:
            clauses = {}
        self._func_def = None
        self._parallel = None
        self._clauses = clauses
        self.device = "mpi"
        self._backend_name = 'MPI'
        self.source_storage = source_storage
        super(AbstractMPIMutator, self, ).__init__()

    def processLlcCodeUnit(self, variables, omp_cl, ast, translate):
        code = ""
        llc_clause_dict = self._get_dict_from_llc_clauses(variables, self.ast)

        # Has directive rnc_result?
        rncresult = False
        rncresult_params = list()
        # Has directive nc_result?
        ncresult = False
        ncresult_params = list()
        # Has directive result?
        result_params = list()
        if llc_clause_dict.has_key('RESULT'):
            if llc_clause_dict['RESULT']:
                for i in variables.result_list:
                    result_params.append([i])
        # Has directive lastresult?
        lastresult = False
        lastresult_params = list()
        # Has directive reduce?
        reduce = False
        reduce_params = list()
        # Has clause reduction?
        reduction_params = self._clauses['REDUCTION']
        # Has clause lastprivate?
        lastprivate_params = self._clauses['LASTPRIVATE']
    
        if (len(ncresult_params) > 0):
           pass

        #--------- INIT ------------# 
        code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_init_unit.tmpl', 'r').read()
        #--------- EXEC ------------# 
        code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_exec_unit.tmpl', 'r').read()
        #--------- COMM ------------# 
        if ((len(result_params) == 1) or ((len(result_params) <= 1) and (len(ncresult_params) > 0))) and (len(rncresult_params) == 0) and (len(lastresult_params) == 0) and (len(lastprivate_params) == 0):
            translate['forall_comm_res_init'] = self.calcResultBufferSize(rncresult_params, result_params, lastresult_params, lastprivate_params)
            translate['forall_comm_res_pack'] = self.packResultData(rncresult_params, result_params, lastresult_params, lastprivate_params)
            translate['forall_comm_res_unpack'] = self.unpackResultData(rncresult_params, result_params, lastresult_params, lastprivate_params)
            if (len(ncresult_params) > 0):
                if (len(result_params) ==  1):
                    translate['var_result_name'] = str(result_params[0][0]).partition(',')[0]
                    translate['var_result_size'] = result_params[0][0].size
                else:
                    translate['var_result_name'] = 'llc_paux'
                    translate['var_result_size'] = '0'
                code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_comm_nc_res_unit.tmpl', 'r').read()
            else:
                translate['var_result_name'] = str(result_params[0][0]).partition(',')[0]
                translate['var_result_size'] = result_params[0][0].size
                code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_comm_unit.tmpl', 'r').read()
        elif ((len(result_params) > 1) or (len(rncresult_params) > 0) or (len(lastresult_params) > 0) or (len(lastprivate_params) > 0)):
            translate['forall_comm_res_init'] = self.calcResultBufferSize(rncresult_params, result_params, lastresult_params, lastprivate_params)
            translate['forall_comm_res_pack'] = self.packResultData(rncresult_params, result_params, lastresult_params, lastprivate_params)
            translate['forall_comm_res_unpack'] = self.unpackResultData(rncresult_params, result_params, lastresult_params, lastprivate_params)
            code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_comm_res_unit.tmpl', 'r').read()
        # Reducciones
        if (len(reduction_params) > 0):
            translate['forall_comm_red_init'] = self.calcReduceBufferSize(reduction_params)
            translate['forall_comm_red_pack_master'] = self.packReduceDataMaster(reduction_params)
            translate['forall_comm_red_unpack_op'] = self.unpackProcessReductionOp(omp_cl, reduction_params)
            translate['forall_comm_red_pack_slave'] = self.packReduceDataSlave(reduction_params)
            translate['forall_comm_red_unpack_slave'] = self.unpackReduceDataSlave(reduction_params)
            code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_comm_red.tmpl', 'r').read()
        #--------- EXIT ------------# 
        code += open(config.WORKDIR + '/Backends/MPI/Templates/forall/forall_exit.tmpl', 'r').read()
        return code

    def calcResultBufferSize(self, rncresult_params, result_params, lastresult_params, lastprivate_params):                                               
        """ Builds the dynamical template for storing the initialization of data. Used in "for" communications.
            
            :param variables: list of llc clauses
            :param omp_cl: list of omp clauses
            :param ast: tree's root
            :result: string built for changing the template for
        """
        cadena = ""
        if (len(result_params) == 0) and (len(rncresult_params) == 0) and (len(lastresult_params) == 0):  # No sea null, que no este
            return cadena
        for decl in result_params:
            cadena += "llc_size_tmp += (" + str(decl[0].size) + ") * sizeof ((" + str(decl[0]).partition(',')[0] + ")[0]);\n";
        # List of lasprivate in omp clause
        for decl in lastprivate_params:
            cadena += "llc_size_tmp += sizeof (" + str(decl.type) + ");\n"
        return cadena

    def packResultData(self, rncresult_params, result_params, lastresult_params, lastprivate_params):
        """
        """
        cadena = ""
        if (len(result_params) == 0) and (len(rncresult_params) == 0) and (len(lastresult_params) == 0):
            return cadena
        # Recorremos la lista de RESULT y para cada variable hacemos:
        for result in result_params:
            cadena += "/* Packing " + str(result[0]).partition(',')[0] + "... */\n"
            cadena += "memcpy (llc_buf_ptr, " + str(result[0]).partition(',')[0] + ", (" + str(result[0].size) + ") * sizeof ((" + str(result[0]).partition(',')[0] + ")[0]));\n"
            cadena += "llc_buf_ptr += (" + str(result[0].size) + ") * sizeof ((" + str(result[0]).partition(',')[0] + ")[0]);\n"
        if (len(lastresult_params) > 0) or (len(lastprivate_params) > 0): 
            cadena += "if (LLC_NAME == (LLC_NUMPROCESSORS - 1)) { \n"
            cadena += "    /* Last execution. Pack variable(s) lastprivate/lastresult for sending */\n\n"
            for decl in lastprivate_params:
                cadena += "\t    /* Packing " + str(decl.name) +"... */\n"
                cadena += "\tmemcpy (llc_buf_ptr, &" + str(decl.name)  + ", sizeof (" + str(decl.name) + "));\n"
                cadena += "\tllc_buf_ptr += sizeof (" + str(decl.name) + ");\n"
            cadena += "}\nelse {\n"
            cadena += "       /* Not last execution. Update llc_buf_ptr is needed (lastprivate/lastresult)*/\n\n"
            for decl in lastprivate_params:
                cadena += "\tllc_buf_ptr += sizeof (" + str(decl.name) + ");\n"
            cadena += "}\n\n"
        return cadena

    def unpackResultData(self, rncresult_params, result_params, lastresult_params, lastprivate_params):
        """
        """
        cadena = ""
        if (len(result_params) == 0) and (len(rncresult_params) == 0) and (len(lastresult_params) == 0):  # No sea null, que no este
            return cadena
        for decl in result_params:
            cadena += "    /* Unpacking " + str(decl[0]).partition(',')[0] + "... */\n"
            cadena += "memcpy (" + str(decl[0]).partition(',')[0] + ", llc_buf_ptr, (" + str(decl[0].size) + ") * sizeof ((" + str(decl[0]).partition(',')[0] + ")[0]));\n"
            cadena += "llc_buf_ptr += (" + str(decl[0].size) + ") * sizeof ((" + str(decl[0]).partition(',')[0] + ")[0]);\n\n"
        if (len(lastprivate_params) > 0) or (len(lastresult_params) > 0): 
            cadena += "if (LLC_NAME != (LLC_NUMPROCESSORS - 1)) {\n"
            cadena += "         /* Not last execution. Unpack variable(s) lastprivate/lastresult received */\n\n"
            for decl in lastprivate_params:
                cadena += "\t    /* Unpacking " + str(decl.name) + "... */\n"
                cadena += "\tmemcpy (&" + str(decl.name) + ", llc_buf_ptr, sizeof (" + str(decl.name) + "));\n"
                cadena += "\tllc_buf_ptr += sizeof (" + str(decl.name) + ");\n,\n"
            cadena += "\n}\nelse {\n"
            cadena += "         /* Last execution. Update llc_buf_ptr is needed (lastprivate/lastresult) */\n\n"
            for decl in lastprivate_params:
                cadena += "\tllc_buf_ptr += sizeof (" + str(decl.name) + "%s);\n"
            cadena += "}\n\n"
        return cadena

    def calcReduceBufferSize(self, reduction_params):
        """
        """
        cadena = ""
        if len(reduction_params) == 0:
            return cadena
        for decl in reduction_params:
           cadena += "llc_buf_size += sizeof (" + str(decl.type) + ");\n"
        return cadena;

    def packReduceDataMaster(self, reduction_params):
        """
        """
        cadena = ""
        if len(reduction_params) == 0:
            return cadena
        cadena += "llc_buf_ptr = llc_buf;\n\n"
        for decl in reduction_params: 
            cadena += "memcpy (llc_buf_ptr, &" + str(decl.name) + ", sizeof(" + str(decl.type) + "));\n"
            cadena += "llc_buf_ptr += sizeof(" + str(decl.type) + ");\n"
        return cadena;
        
    def _operator_reduction(self, clause_type, decl):
        switch = { 
                    '+': "" + str(decl.name) + " += (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '*': "" + str(decl.name) + " *= (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '-': "" + str(decl.name) + " += (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '&': "" + str(decl.name) + " &= (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '|': "" + str(decl.name) + " |= (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '^': "" + str(decl.name) + " ^= (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '&&': "" + str(decl.name) + " = " + str(decl.name) + " && (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                    '||': "" + str(decl.name) + " = " + str(decl.name) + " || (*(" + str(decl.type) + " *) llc_buf_ptr);\n",
                }
        return switch[clause_type]

    
    def unpackProcessReductionOp(self, omp_cl, reduction_params):
        """
        """
        cadena = ""
        if len(reduction_params) == 0:
            return cadena
        cadena += "llc_buf_ptr = llc_buf;\n\n"
        if len(reduction_params) == 0:
            return cadena
        for clause in omp_cl:
            if clause.name == 'REDUCTION':
                for decl in reduction_params:
                    try:
                        cadena += self._operator_reduction(clause.type, decl);
                        cadena += "llc_buf_ptr +=  sizeof(" + str(decl.type) + ");\n"
                    except KeyError:
                        print "KEY ERROR: MIN/MAX O ERROR!!"
                        raise AbortMutationException("ERROR: OMP reduction operator not valid (" + clause.type + ")")
        return cadena

    def packReduceDataSlave(self, reduction_params):
        """
        """
        cadena = ""
        if len(reduction_params) == 0:
            return cadena
        cadena += "llc_buf_ptr = llc_buf;\n\n"
        for decl in reduction_params: 
            cadena += "memcpy (llc_buf_ptr, &" + str(decl.name) + ", sizeof(" + str(decl.type) + "));\n"
            cadena += "llc_buf_ptr += sizeof(" + str(decl.type) + ");\n"
        return cadena

    def unpackReduceDataSlave(self, reduction_params):
        """
        """
        cadena = ""
        if len(reduction_params) == 0:
            return cadena
        cadena += "llc_buf_ptr = llc_buf;\n\n"
        for decl in reduction_params: 
            cadena += "memcpy (&" + str(decl.name) + ", llc_buf_ptr, sizeof(" + str(decl.type) + "));\n"
            cadena += "llc_buf_ptr += sizeof(" + str(decl.type) + ");\n"
        return cadena;

    def _get_dict_from_llc_clauses(self, llc_clauses, ast):
        """
             Return a dict of clauses from a list of OmpClause objects
            
              Example: [OmpClause('RESULT', ...), OmpClause('NC_RESULT', ...)]
                 will return:  {'RESULT' : [....] , 'NC_RESULT' : [...]}

            :return: dict with llc clauses
        """
        llc_clause_names = ['RESULT', 'NC_RESULT']
        llc_clause_dict = {}
        for name in llc_clause_names:
            llc_clause_dict[name] = False
        # The FOR clause hasn't got attribute name
        try:
            clause_name = llc_clauses.name
            if clause_name in llc_clause_names:
                llc_clause_dict[clause_name] = True
        except AttributeError:
            return llc_clause_dict
        return  llc_clause_dict

def extract_decls_loop_declaration(objetive, ast):
    loop = objetive.parent
    try:
        if loop.name in ['RESULT', 'NC_RESULT']:
            loop = loop.stmt
    except:
        pass
    loop_parameters = OmpParametrize().apply(loop,)
    decl_list = []
    for i in loop_parameters['global_vars']:
        decl_list.append(i.node)
    return decl_list

def mpi_outline_loop_stmts(objetive, ast, file_name = 'fake_function', template_source = None):
    """
        Outline a mpi function from a For Loop
        :param objetive: statements from a loop (For.stmt)
    """
    # assert block_is_SESE(objetive)
    # Do not replace original subtree
    om = OutlineMutator(name = file_name + "_support", parameter_style = ParameterStyle.C_POINTER, change_code_block = False)
    om.add_additional_decls(extract_decls_loop_declaration(objetive, ast))
    om.apply(objetive)
    forall_template = """
    <%!
        from yacf.Backends.Common.TemplateEngine.Functions import is_complex_type
    %>

    %for typedef in typedeflist:
        ${typedef};
    %endfor
     
    ${outlinedProto};
    ${outlinedProto} {
        %for var in parameters:
            %if not is_complex_type(var.ptr):
                ${str(var.declaration) + str(var) + " = " +  "*_" + str(var)};
            %endif
        %endfor   
        ${template_code}
        %for var in parameters:
            %if not is_complex_type(var.ptr):
                ${"*_" + str(var) + " = " + str(var)};
            %endif
        %endfor   
    }

    """
    subs_dir = {'typedeflist' : om.getTypedefList(), 
                'outlinedProto' : om.getOutlinedPrototype(), 
                'outlinedFunction' : om.getOutlinedFunction(),
                'template_code' : str(template_source),
                'parameters' : om.getParameterList(),
                'ArrayDecl' : str(c_ast.ArrayDecl),
                'PtrDecl' : str(c_ast.PtrDecl)}
    forall_source = AbstractMutator().fillTemplate(forall_template, subs_dir)
    # Replace the Node
    Forall_pragma = objetive.parent.parent.parent 
    new = c_ast.Compound(block_items = [om.getOutlinedFunctionCall()])
    ReplaceTool(new_node = new, old_node = Forall_pragma).apply(Forall_pragma.parent, 'block_items')
    return forall_source
